# -*- coding: utf-8 -*-

import pickle
import numpy as np
import pandas as pd

from StoreBilibiliData.CreateTable import uploader, video, comments, dm

from sqlalchemy import create_engine, select
from sqlalchemy.sql.selectable import Select
from sqlalchemy.engine.base import Engine, Connection
from sqlalchemy.engine.result import ResultProxy, RowProxy


def load_uploader_data_to_pandas(*database_names: [str]) -> pd.DataFrame:
    ser = []
    for i in database_names:
        eng = create_engine('sqlite:///' + i, echo=False, encoding='utf-8')  # type: Engine
        con = eng.connect()  # type: Connection

        sel = select([uploader.c.uid, uploader.c.name])  # type: Select
        rp = con.execute(sel)  # type: ResultProxy

        for r in rp:  # type: RowProxy
            tmp = {
                'uid': r.uid,
                'name': r.name,
            }
            ser.append(pd.Series(tmp))

    return pd.DataFrame(ser).drop_duplicates(subset='uid').set_index(keys=['uid'])


def load_video_data_to_pandas(*database_names: [str]) -> pd.DataFrame:
    ser = []
    for i in database_names:
        eng = create_engine('sqlite:///' + i, echo=False, encoding='utf-8')  # type: Engine
        con = eng.connect()  # type: Connection

        sel = select([video])  # type: Select
        rp = con.execute(sel)  # type: ResultProxy

        for r in rp:  # type: RowProxy
            tmp = {
                'av': r.av,
                'bv': r.bv,
                'comment_count': r.comment_count,
                'play_count': r.play_count,
                'title': r.title,
                'description': r.description,
                'uploader_id': r.uploader_id,
                'upload_time': r.upload_time,
            }
            ser.append(pd.Series(tmp))

    return pd.DataFrame(ser).set_index(keys=['av'])


def load_dm_data_to_pandas(*database_names: [str]) -> pd.DataFrame:
    ser = []
    for i in database_names:
        eng = create_engine('sqlite:///' + i, echo=False, encoding='utf-8')  # type: Engine
        con = eng.connect()  # type: Connection

        sel = select([dm])  # type: Select
        rp = con.execute(sel)  # type: ResultProxy

        for r in rp:  # type: RowProxy
            tmp = {
                'video_id': r.video_id,
                'content': r.content,
            }
            ser.append(pd.Series(tmp))

    df = pd.DataFrame(ser).set_index(keys=['video_id'])
    df['content'] = df['content'].str.replace(' ', '')
    return df


def load_comments_data_to_pandas(*database_names: [str]) -> pd.DataFrame:
    ser = []
    for i in database_names:
        eng = create_engine('sqlite:///' + i, echo=False, encoding='utf-8')  # type: Engine
        con = eng.connect()  # type: Connection

        sel = select([comments])  # type: Select
        rp = con.execute(sel)  # type: ResultProxy

        for r in rp:  # type: RowProxy
            tmp = {
                'rp_id': r.rp_id,
                'video_id': r.video_id,
                'likes': r.likes,
                'root_comment': r.root_comment,
                'content': r.content,
                'comment_time': r.comment_time,
            }
            ser.append(pd.Series(tmp))

    df = pd.DataFrame(ser).set_index(keys=['rp_id'])  # type: pd.DataFrame
    df['content'] = df['content'].str.replace(' ', '')

    def __split(x: pd.Series) -> pd.Series:
        x['content'] = x['content'].split(':', 1)[-1]
        return x

    df = df.apply(func=__split, axis=1)
    return df


if __name__ == '__main__':
    databases = [
        '../DatabaseFile/bilibili 2020-01.db',
        '../DatabaseFile/bilibili 2020-02.db',
        '../DatabaseFile/bilibili 2020-03.db',
        '../DatabaseFile/bilibili 2020-04.db',
        '../DatabaseFile/bilibili 2020-05.db',
    ]

    ups = load_uploader_data_to_pandas(*databases)
    print(ups)

    vid = load_video_data_to_pandas(*databases)
    print(vid)

    dms = load_dm_data_to_pandas(*databases)
    print(dms)

    cts = load_comments_data_to_pandas(*databases)
    print(cts)

    pickle.dump(obj=ups, file=open(file='../PickledFile/uploader.pickled', mode='wb'))
    pickle.dump(obj=vid, file=open(file='../PickledFile/video.pickled', mode='wb'))
    pickle.dump(obj=dms, file=open(file='../PickledFile/dm.pickled', mode='wb'))
    pickle.dump(obj=cts, file=open(file='../PickledFile/comments.pickled', mode='wb'))
